// TODO:
//  1. 对于 pool 管理的堆分配数据，需要检查 OOM
//  2. 检查 u64/usize 转换的合法性 （32位系统优化）
use std::{borrow::Cow, sync::Arc};

use crate::utils::TonglunBitcast;
use parking_lot::lock_api::ReentrantMutex;

use crate::{
    boxes::{ArrayViewBox, Complex64Box, StringBox},
    error::TonglunError,
    functions::TonglunFunction,
    globals::{
        POOL_ARRAYS, POOL_COMPLEX64, POOL_ERRORS, POOL_RUNTIME, POOL_STRINGS, POOL_TUPLES,
        REGISTERED_FUNCTIONS, REGISTERED_RUNTIMES,
    },
    pools::PoolEntry,
    runtime::Runtime,
    utils::{FFIUtils, NullStr},
    values::{ArrayDType, Complex32, RustValue, ValueRef, ValueTag},
};

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_runtime(
    // [OUT]
    out_runtime_id: *mut u32,

    // [IN]
    dlhandle: *const std::ffi::c_void,
    name: *const std::ffi::c_char,
    thread_safe: u8,
    string_drop_callback: *const std::ffi::c_void,
    array_drop_callback: *const std::ffi::c_void,
) -> u64 {
    let thread_safe = thread_safe != 0;
    let mut registration_lock = REGISTERED_RUNTIMES.write();
    let dlhandle_as_usize = dlhandle as usize;
    let name = match FFIUtils::c_str_to_string(name) {
        Ok(s) => s,
        Err(e) => {
            return e.into_usize() as u64;
        }
    };
    if registration_lock.contains_key(&dlhandle_as_usize) {
        let error = TonglunError::RuntimeAlreadyRegistered(dlhandle_as_usize, name);
        return error.into_usize() as u64;
    }
    let mut runtime_lock = POOL_RUNTIME.write();
    let runtime = runtime_lock.alloc_with_handle(Runtime {
        id: 0,
        name: NullStr::from(name),
        thread_safe,
        string_drop_callback: if thread_safe {
            None
        } else {
            Some(unsafe { std::mem::transmute(string_drop_callback) })
        },
        array_drop_callback: if thread_safe {
            None
        } else {
            Some(unsafe { std::mem::transmute(array_drop_callback) })
        },
        thread_defer_dealloc_strings: boxcar::Vec::new(),
        thread_defer_dealloc_arrays: boxcar::Vec::new(),
        lock: Arc::new(ReentrantMutex::new(())),
    });
    let runtime_id = runtime.to_usize() as u32;
    registration_lock.insert(dlhandle_as_usize, runtime_id);
    unsafe { *out_runtime_id = runtime_id };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_free_value(value_handle: *const ValueRef) -> u64 {
    let type_safe_value = RustValue::from_value_ref(unsafe { *value_handle });
    match type_safe_value {
        RustValue::None => {}
        RustValue::Bool(_) => {}
        RustValue::Int64(_) => {}
        RustValue::Float64(_) => {}
        RustValue::Complex32(_) => {}
        RustValue::Complex64(pool_ref) => {
            let queue = POOL_COMPLEX64
                .write()
                .free_queue(PoolEntry::new(pool_ref.to_usize()));

            drop(queue);
        }
        RustValue::String(pool_ref) => {
            let queue = POOL_STRINGS
                .write()
                .free_queue(PoolEntry::new(pool_ref.to_usize()));

            drop(queue);
        }
        RustValue::Array(pool_ref) => {
            let queue = POOL_ARRAYS
                .write()
                .free_queue(PoolEntry::new(pool_ref.to_usize()));

            drop(queue);
        }
        RustValue::HList(pool_ref) => {
            let queue = POOL_TUPLES
                .write()
                .free_queue(PoolEntry::new(pool_ref.to_usize()));

            drop(queue);
        }
    }
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_string_copied(
    out: *mut ValueRef,
    input: *const std::ffi::c_char,
) -> u64 {
    let input = match FFIUtils::c_str_to_string(input) {
        Ok(s) => s,
        Err(e) => {
            return e.into_usize() as u64;
        }
    };
    let mut string_lock = POOL_STRINGS.write();
    let entry = string_lock.alloc_with_handle(StringBox {
        handle: Default::default(),
        length: input.len() as u64,
        data: Cow::Owned(input),
        foreign_ob: std::ptr::null_mut(),
        runtime_id: 0,
    });
    let value_ref = RustValue::String(entry.to_pool_ref()).to_value_ref();
    unsafe { *out = value_ref };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_string_borrowed(
    out: *mut ValueRef,
    runtime_id: u32,
    input: *const std::ffi::c_char,
    foreign_ob: *mut std::ffi::c_void,
) -> u64 {
    let n = FFIUtils::c_strlen(input);
    let v = std::slice::from_raw_parts(input as *const u8, n);
    let s = match std::str::from_utf8(v) {
        Ok(s) => s,
        Err(_) => {
            return TonglunError::InvalidUtf8String.into_usize() as u64;
        }
    };
    let mut string_lock = POOL_STRINGS.write();
    let entry = string_lock.alloc_with_handle(StringBox {
        handle: Default::default(),
        length: n as u64,
        data: Cow::Borrowed(s),
        foreign_ob,
        runtime_id,
    });
    let value_ref = RustValue::String(entry.to_pool_ref()).to_value_ref();
    unsafe { *out = value_ref };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_none(out: *mut ValueRef) -> u64 {
    let value_ref = RustValue::None.to_value_ref();
    unsafe { *out = value_ref };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_bool(out: *mut ValueRef, value: u8) -> u64 {
    let value = value != 0;
    let value_ref = RustValue::Bool(value).to_value_ref();
    unsafe { *out = value_ref };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_int64(out: *mut ValueRef, value: i64) -> u64 {
    let value_ref = RustValue::Int64(value).to_value_ref();
    unsafe { *out = value_ref };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_float64(out: *mut ValueRef, value: f64) -> u64 {
    let value_ref = RustValue::Float64(value).to_value_ref();
    unsafe { *out = value_ref };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_complex32(out: *mut ValueRef, real: f32, imag: f32) -> u64 {
    let complex_box = Complex32 { real, imag };
    let value_ref = RustValue::Complex32(complex_box).to_value_ref();
    unsafe { *out = value_ref };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_complex64(out: *mut ValueRef, real: f64, imag: f64) -> u64 {
    let complex_box = Complex64Box { real, imag };
    let mut complex_lock = POOL_COMPLEX64.write();
    let entry = complex_lock.alloc(complex_box);
    let value_ref = RustValue::Complex64(entry.to_pool_ref()).to_value_ref();
    unsafe { *out = value_ref };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_array(
    // [OUT]
    out: *mut ValueRef,

    // [IN]
    runtime_id: u32,
    typecode: *const std::ffi::c_char,
    data: *const std::ffi::c_void,
    ndims: i64,
    dimensions: *const usize,
    strides: *const isize,
    foreign_ob: *mut std::ffi::c_void,
    readonly: u8,
    fortran_order: u8,
) -> u64 {
    let readonly = readonly != 0;
    let str_len = FFIUtils::c_strlen(typecode);
    let typecode = std::slice::from_raw_parts(typecode as *const u8, str_len);

    let slice_dimensions: &[usize] = std::slice::from_raw_parts(dimensions, ndims as usize);
    let slice_strides: &[isize] = std::slice::from_raw_parts(strides, ndims as usize);

    let dtype = match ArrayDType::from_typecode(typecode) {
        Ok(dtype) => dtype,
        Err(e) => {
            return e.into_usize() as u64;
        }
    };

    let mut array_lock = POOL_ARRAYS.write();
    let entry = array_lock.alloc_with_handle(ArrayViewBox {
        handle: Default::default(),
        data: data,
        dimensions: slice_dimensions.to_vec(),
        strides: slice_strides.to_vec(),
        ndims: ndims as u64,
        dtype,
        foreign_ob,
        runtime_id,
        readonly,
        fortran_order: fortran_order != 0,
    });

    let value_ref = RustValue::Array(entry.to_pool_ref()).to_value_ref();
    unsafe { *out = value_ref };
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_create_tuple(
    out: *mut ValueRef,
    nitems: u64,
    items: *const ValueRef,
) -> u64 {
    let mut vec = std::vec::Vec::new();
    if let Err(_) = vec.try_reserve_exact(nitems as usize) {
        return TonglunError::OutOfMemory.into_usize() as u64;
    }

    for i in 0..nitems {
        let item = unsafe { *items.add(i as usize) };
        vec.push(item);
    }

    let entry = {
        let mut pool_lock = POOL_TUPLES.write();
        pool_lock.alloc(vec)
    };

    let value_ref = RustValue::HList(entry.to_pool_ref()).to_value_ref();
    unsafe { *out = value_ref };

    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_value_tag(value: *const ValueRef) -> ValueTag {
    (*value).tag
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_get_bool(value: *const ValueRef) -> u8 {
    let value = *value;
    if value.data != 0 {
        1
    } else {
        0
    }
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_get_int64(value: *const ValueRef) -> i64 {
    let value = *value;
    i64::_tl_from_unsigned(value.data)
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_get_float64(value: *const ValueRef) -> f64 {
    let value = *value;
    f64::_tl_from_unsigned(value.data)
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_load_complex32(
    real: *mut f32,
    imag: *mut f32,
    value: *const ValueRef,
) {
    let value = *value;
    unsafe { *real = Complex32::_tl_from_unsigned(value.data).real };
    unsafe { *imag = Complex32::_tl_from_unsigned(value.data).imag };
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_load_complex64(
    real: *mut f64,
    imag: *mut f64,
    value: *const ValueRef,
) {
    let handle = (*value).data;
    let pool_lock = POOL_COMPLEX64.read();
    let Some(entry) = pool_lock.get(PoolEntry::new(handle as usize)) else {
        return;
    };
    unsafe { *real = entry.real };
    unsafe { *imag = entry.imag };
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_load_string(
    out: *mut *const std::ffi::c_char,
    value: *const ValueRef,
) {
    let handle = (*value).data;
    let pool_lock = POOL_STRINGS.read();
    let Some(entry) = pool_lock.get(PoolEntry::new(handle as usize)) else {
        return;
    };
    unsafe { *out = entry.data.as_ptr() as *const std::ffi::c_char };
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_get_tuple_length(value: *const ValueRef) -> u64 {
    let handle = (*value).data;
    let pool_lock = POOL_TUPLES.read();
    let Some(entry) = pool_lock.get(PoolEntry::new(handle as usize)) else {
        return 0;
    };
    entry.len() as u64
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_take_tuple(out: *mut ValueRef, value: *const ValueRef) {
    let handle = (*value).data;
    let mut pool_lock = POOL_TUPLES.write();
    let Some(elements) = pool_lock.get(PoolEntry::new(handle as usize)) else {
        return;
    };
    std::ptr::copy(elements.as_ptr(), out as *mut ValueRef, elements.len());
    pool_lock.free(PoolEntry::new(handle as usize));
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_load_array_info(
    // OUT
    data: *mut *const std::ffi::c_void,
    dimensions: *mut *const usize,
    strides: *mut *const isize,
    ndims: *mut i64,
    dtype: *mut ArrayDType,
    readonly: *mut u8,
    fortran_order: *mut u8,
    // IN
    handle: u64,
) {
    let pool_lock = POOL_ARRAYS.read();
    let Some(entry) = pool_lock.get(PoolEntry::new(handle as usize)) else {
        return;
    };
    unsafe { *data = entry.data };
    unsafe { *dimensions = entry.dimensions.as_ptr() };
    unsafe { *strides = entry.strides.as_ptr() };
    unsafe { *ndims = entry.ndims as i64 };
    unsafe { *dtype = entry.dtype };
    unsafe { *readonly = entry.readonly as u8 };
    unsafe { *fortran_order = entry.fortran_order as u8 };
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_register_function(
    runtime_id: u32,
    name: *const std::ffi::c_char,
    fptr: *const std::ffi::c_void,
    data: u64,
) -> u64 {
    let name = match FFIUtils::c_str_to_string(name) {
        Ok(s) => s,
        Err(e) => {
            return e.into_usize() as u64;
        }
    };
    let thread_safe = {
        let runtime_lock = POOL_RUNTIME.read();
        let Some(runtime) = runtime_lock.get(PoolEntry::new(runtime_id as usize)) else {
            return TonglunError::RuntimeNotFound(runtime_id as usize).into_usize() as u64;
        };
        runtime.thread_safe
    };

    let name_without_null_terminator = name.clone();

    let function = TonglunFunction {
        fptr: unsafe { std::mem::transmute(fptr) },
        data,
        runtime_id,
        name: Arc::new(NullStr::from(name.clone())),
        thread_safe,
    };

    // // check if the function is already registered
    // if REGISTERED_FUNCTIONS.contains_key(&name_without_null_terminator) {
    //     return TonglunError::FunctionAlreadyRegistered(name_without_null_terminator).into_usize()
    //         as u64;
    // }

    REGISTERED_FUNCTIONS.insert(name_without_null_terminator, function);
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_call_function(
    out: *mut ValueRef,
    name: *const std::ffi::c_char,
    args: *const ValueRef,
    nargs: u64,
) -> u64 {
    let strlen = FFIUtils::c_strlen(name);
    let name = std::slice::from_raw_parts(name as *const u8, strlen);
    let name = match std::str::from_utf8(name) {
        Ok(s) => s,
        Err(_) => {
            return TonglunError::InvalidUtf8String.into_usize() as u64;
        }
    };

    let function: Option<TonglunFunction> = REGISTERED_FUNCTIONS.get(name).map(|it| it.clone());

    let Some(function) = function else {
        return TonglunError::FunctionNotFound(name.to_string()).into_usize() as u64;
    };

    function.do_call(out, args, nargs)
}

/// The buffer capacity should be buffer_len_minus_1 + 1;
#[no_mangle]
pub unsafe extern "C" fn tonglun_try_load_error_message(
    // [OUT]
    buffer: *mut std::ffi::c_char,

    // [In],
    buffer_len_minus_1: u64,
    error_handle: u64,
) -> u64 {
    let error_string = {
        let lock = POOL_ERRORS.read();
        let Some(entry) = lock.get(PoolEntry::new(error_handle as usize)) else {
            return TonglunError::ErrorNotFound(error_handle as usize).into_usize() as u64;
        };
        entry.to_string()
    };
    let error_string_len = error_string.len();
    if error_string_len > buffer_len_minus_1 as usize {
        return TonglunError::NotEnoughBufferLength(error_string_len).into_usize() as u64;
    }

    // otherwise, fill the buffer with the error string
    let error_string_ptr = error_string.as_ptr() as *mut u8;
    std::ptr::copy(error_string_ptr, buffer as *mut u8, error_string_len);
    // write the null terminator
    unsafe {
        *buffer.add(error_string_len) = 0;
    }

    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_error_buflen(
    // [IN]
    error_handle: u64,
) -> u64 {
    let expected_buffer_len = {
        let lock = POOL_ERRORS.read();
        let Some(TonglunError::NotEnoughBufferLength(expected_buffer_len)) =
            lock.get(PoolEntry::new(error_handle as usize))
        else {
            return 0;
        };
        *expected_buffer_len
    };
    expected_buffer_len as u64
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_free_error(error_handle: u64) -> u64 {
    let mut error_lock = POOL_ERRORS.write();
    error_lock.free(PoolEntry::new(error_handle as usize));
    0
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_binding_error(
    language: *const std::ffi::c_char,
    message: *const std::ffi::c_char,
) -> u64 {
    let language = match FFIUtils::c_str_to_string(language) {
        Ok(s) => s,
        Err(e) => {
            // TODO: print the real error message
            println!("OOM triggerd when handling binding error");
            return e.into_usize() as u64;
        }
    };
    let message = match FFIUtils::c_str_to_string(message) {
        Ok(s) => s,
        Err(e) => {
            // TODO: print the real error message
            println!("OOM triggerd when handling binding error");
            return e.into_usize() as u64;
        }
    };
    let error = TonglunError::BindingError { language, message };
    error.into_usize() as u64
}

#[no_mangle]
pub unsafe extern "C" fn tonglun_run_deferred_free(runtime_id: u32) -> u64 {
    let mut runtime_lock = POOL_RUNTIME.write();
    let Some(runtime) = runtime_lock.get_mut(PoolEntry::new(runtime_id as usize)) else {
        return TonglunError::RuntimeNotFound(runtime_id as usize).into_usize() as u64;
    };

    if let Some(string_drop_callback) = runtime.string_drop_callback {
        for item in runtime.thread_defer_dealloc_arrays.iter() {
            string_drop_callback(*item.1);
        }
        runtime.thread_defer_dealloc_strings.clear();
    }

    if let Some(array_drop_callback) = runtime.array_drop_callback {
        for item in runtime.thread_defer_dealloc_arrays.iter() {
            array_drop_callback(*item.1);
        }
        runtime.thread_defer_dealloc_arrays.clear();
    }
    0
}
